#encoding=utf-8
from flask import Flask, request, send_from_directory, render_template, redirect, url_for, flash
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, SubmitField
from wtforms.validators import InputRequired, Length
from werkzeug.security import generate_password_hash, check_password_hash
from werkzeug.utils import secure_filename
import os
import math
# Your existing Flask app setup
app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key'  # Replace 'your-secret-key' with a real secret key
app.config['UPLOAD_FOLDER'] = 'uploaded_files'
if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

# Flask-Login setup
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'

# Mock database of users
users = {'admin': {'password': generate_password_hash('admin')}}

# User session management setup
class User(UserMixin):
    pass

@login_manager.user_loader
def user_loader(username):
    if username not in users:
        return

    user = User()
    user.id = username
    return user

# Forms
class LoginForm(FlaskForm):
    username = StringField('Username', validators=[InputRequired(), Length(min=4, max=15)])
    password = PasswordField('Password', validators=[InputRequired(), Length(min=3, max=80)])
    submit = SubmitField('Login')

# Authentication routes
@app.route('/login', methods=['GET', 'POST'])
def login():
    form = LoginForm()

    if form.validate_on_submit():
        user = users.get(form.username.data)
        if user and check_password_hash(user['password'], form.password.data):
            user_obj = User()
            user_obj.id = form.username.data
            login_user(user_obj)
            return redirect(url_for('upload_file'))

        flash('Invalid username or password')

    return render_template('login.html', form=form)

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('login'))


def convert_size(size_bytes):
    if size_bytes == 0:
        return "0B"
    size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
    i = int(math.floor(math.log(size_bytes, 1024)))
    p = math.pow(1024, i)
    s = round(size_bytes / p, 2)
    return f"{s} {size_name[i]}"

# Protect the upload routes
@app.route('/')
@login_required
def upload_file():
    files = os.listdir(app.config['UPLOAD_FOLDER'])
    files_with_sizes = []
    for file in files:
        path = os.path.join(app.config['UPLOAD_FOLDER'], file)
        size = os.path.getsize(path)
        readable_size = convert_size(size)
        files_with_sizes.append((file, readable_size))
    return render_template('upload.html', files=files_with_sizes)

@app.route('/upload', methods=['POST'])
@login_required
def upload_file_post():
    if 'file' not in request.files:
        flash('没有文件部分')
        return redirect(request.url)
    file = request.files['file']
    if file.filename == '':
        flash('没有选中文件')
        return redirect(request.url)
    if file:
        filename = secure_filename(file.filename)
        save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(save_path)
        flash('文件成功上传')
        return redirect(url_for('upload_file'))



@app.route('/uploads/<filename>')
@login_required
def uploaded_file(filename):

    print("test size")
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

@app.route('/delete/<filename>', methods=['POST'])
@login_required
def delete_file(filename):
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if os.path.exists(file_path):
        os.remove(file_path)
        flash('文件删除成功')
    else:
        flash('文件不存在')
    return redirect(url_for('upload_file'))

if __name__ == "__main__":
    app.run(host='0.0.0.0', debug=True, port=80)

